// coding: utf-8
// Time: 2022-08-24 12:22:09
// Desc: 系统调度
// compile: gcc ./sched_test.c -o sched_test -pthread

#define _GNU_SOURCE
#include "stdio.h"
#include "stdlib.h"
#include "sched.h"
#include "time.h"
#include "sys/timeb.h"
#include "sys/sysinfo.h"
#include "unistd.h"
#include "pthread.h"

struct ThreadParam {
    float rate;
    int interval;
    int timeout;
};

long long get_system_time()
{
    struct timespec ts;
    (void)clock_gettime(CLOCK_MONOTONIC, &ts);
    long long milliseconds = (ts.tv_sec * 1000) + (ts.tv_nsec / 1000000);
    return milliseconds;
}

void* exec_on_cpu(void* arg) {
    struct ThreadParam param = *(struct ThreadParam *)arg;
    float rate = param.rate;
    int interval = param.interval;
    int timeout = param.timeout;
    long long t_start = get_system_time();
    while (get_system_time() - t_start < timeout * 1000) {
        int t_exec = get_system_time() - t_start;
        if (t_exec / interval % 1000 < rate * 1000) {
            // pass
        } else {
            usleep(interval * rate * 1000);
        }
    }
    return 0;
}

pthread_t set_cpu_rate_of_cpu(int timeout, int interval, float rate, int cpu_index) {
    if (rate < 0 || rate > 1.0) {
        printf("Input usage rate is %f, which must be a float value between 0 and 1!\n", rate);
        return 0;
    }
    if (interval <= 0) {
        printf("Input interval is %d, which must be a positive number!\n", interval);
        return 0;
    }
    if (timeout <= 0) {
        printf("Input timeout is %d, which must be a positive number!\n", timeout);
        return 0;
    }
    if (cpu_index < 0) {
        printf("Input cpu index is %d, which must be a positive number!\n", cpu_index);
    }
    cpu_set_t mask;
    CPU_ZERO(&mask);  
    CPU_SET(cpu_index, &mask);
    if (sched_setaffinity(0, sizeof(mask), &mask) == -1) {
        printf("Cannot set to cpu %d", cpu_index);
        return 0;
    }
    pthread_t t;
    struct ThreadParam param;
    param.rate = rate;
    param.interval = interval;
    param.timeout = timeout;
    pthread_create(&t, 0, exec_on_cpu, &param);
    return t;
}

void set_cpu_rate_of_system(int timeout, int interval, float rate) {
    if (rate < 0 || rate > 1.0) {
        printf("Input usage rate is %f, which must be a float value between 0 and 1!\n", rate);
        return;
    }
    if (interval <= 0) {
        printf("Input interval is %d, which must be a positive number!\n", interval);
        return;
    }
    if (timeout <= 0) {
        printf("Input timeout is %d, which must be a positive number!\n", timeout);
        return;
    }
    int cpu_num = get_nprocs_conf();
    pthread_t t_vec[cpu_num];
    int i = 0;
    for (; i < cpu_num; ++i) {
        cpu_set_t mask;
        CPU_ZERO(&mask);
        CPU_SET(i, &mask);
        if (sched_setaffinity(0, sizeof(mask), &mask) == -1) {
            printf("Cannot set to cpu %d", i);
            return;
        }
        pthread_t t;
        struct ThreadParam param;
        param.rate = rate;
        param.interval = interval;
        param.timeout = timeout;
        pthread_create(&t, 0, exec_on_cpu, &param);
        t_vec[i] = t;
    }
    for (i = 0; i < cpu_num; ++i) {
        pthread_join(t_vec[i], NULL);
    }
}

void set_cpu_rate_of_task(int timeout, int interval, float rate) {
    set_cpu_rate_of_system(timeout, interval, rate);
}

#if 0
int main(int argc, char** argv) {
    if (argc < 3) {
        printf("usage: ./sched_test timeout interval rate [cpu_index]\n");
        return 0;
    }
    int timeout = atoi(argv[1]);
    int interval = atoi(argv[2]);
    float rate = atof(argv[3]);
    int param4 = 0;
    if (argc == 5) {
        param4 = atoi(argv[4]);  // param4 can be used as cpu_index or process_id
    }

    // Test set_cpu_rate_of_cpu
    // pthread_t t = set_cpu_rate_of_cpu(timeout, interval, rate, param4);
    // if (t == 0) return 1;
    // pthread_join(t, NULL);
    // Test set_cpu_rate_of_system
    set_cpu_rate_of_system(timeout, interval, rate);
    return 0;
}
#endif
